import numpy as np
import matplotlib.pyplot as plt
import os
import random
import cv2
os.chdir(r"C:\Users\cai\Desktop\光电竞赛迷宫寻宝\neural_net_classify")

def loaddatas():#用于分类
    list_1=os.listdir('./traindata/')
    images=[]
    anatations=[]
    for fullname in list_1:
        file_name,file_extend = os.path.splitext(fullname)
        if file_name.isdigit():
            list0=os.listdir('./traindata/'+fullname)
            for image in list0:
                img=cv2.imread('./traindata/'+fullname+'/'+image)
                if img[0][0][0].dtype == 'uint8':
                    div_flag=True
                else:
                    div_flag=False
                #img = plt.imread('./traindata/1/0.jpg')
                img=cv2.resize(img,(30,60))
                images.append(img)
                anatation=np.zeros(4).astype("uint8")
                anatation[int(file_name)-1]=1
                anatations.append(anatation)

    c = list(zip(images, anatations))
    random.shuffle(c)
    images,anatations=zip(*c)
    if div_flag:
        images_ndarray=np.array(images)/255
    anatations_ndarry=np.array(anatations)
    return images_ndarray,anatations_ndarry

def d_loaddatas():#用于辨别
    images=[]
    anatations=[]

    list_1=os.listdir('./d_traindata/diff1')
    for img_name in list_1:
        img1=plt.imread('./d_traindata/diff1/'+img_name)
        img2 = plt.imread('./d_traindata/diff2/' + img_name)
        img1=img1.astype('int16')
        img2 = img2.astype('int16')
        img=(img1-img2)
        images.append(img)
        anatation=0.0
        anatations.append(anatation)

    list_2=os.listdir('./d_traindata/same1')
    for img_name in list_2:
        img1=plt.imread('./d_traindata/same1/'+img_name)
        img2 = plt.imread('./d_traindata/same2/' + img_name)
        img1=img1.astype('int16')
        img2 = img2.astype('int16')
        img=(img1-img2)
        images.append(img)
        anatation=1.0
        anatations.append(anatation)

    c = list(zip(images, anatations))
    random.shuffle(c)
    images,anatations=zip(*c)

    images_ndarray=np.abs(np.array(images)/255)
    anatations_ndarry=np.array(anatations)
    return images_ndarray,anatations_ndarry


if __name__ == "__main__":
    a,b=loaddatas()
    # images_ndarry=np.array(images)/255
    # anatations_ndarry=np.array(anatations)/255
    # return anatations_ndarry,images_ndarry
    pass
